from datetime import datetime, timedelta
from vnpy_ctastrategy import (
    CtaTemplate, StopOrder, TickData, BarData, 
    TradeData, OrderData, BarGenerator, ArrayManager
)
import pandas as pd
import tushare as ts
from typing import Optional


class MarketDataLoader:
    """统一获取宏观经济与市场指数数据"""
    def __init__(self, token: str, 
                 start_date: Optional[datetime] = None,
                 end_date: Optional[datetime] = None):
        """
        :param token: Tushare token
        :param start_date: 数据开始日期（包含）
        :param end_date: 数据结束日期（包含）
        """
        self.pro = ts.pro_api(token)
        self.start_date = start_date or datetime(1990, 1, 1)
        self.end_date = end_date or datetime.now()
        self._init_data()

    def _init_data(self):
        """初始化指定时间范围内的数据"""
        self.gdp = self._load_gdp()
        self.macro_monthly = self._load_macro_monthly()
        self.zz1000 = self._load_zz1000()

    def _date_to_quarter(self, dt: datetime) -> str:
        """将日期转换为Tushare季度格式"""
        year = dt.year
        quarter = (dt.month - 1) // 3 + 1
        return f"{year}Q{quarter}"

    def _date_to_month(self, dt: datetime) -> str:
        """转换为Tushare月度格式"""
        return dt.strftime("%Y%m")

    def _load_gdp(self) -> pd.Series:
        """加载指定时间范围内的GDP数据"""
        # 转换时间参数
        start_q = self._date_to_quarter(self.start_date)
        end_q = self._date_to_quarter(self.end_date)
        
        # 获取原始数据
        df = self.pro.cn_gdp(
            start_q=start_q,
            end_q=end_q,
            fields='quarter,gdp_yoy'
        )
        
        # 日期处理
        df['date'] = pd.to_datetime(
            df['quarter'].str[:4] + 
            df['quarter'].str[-1].map({'1':'03','2':'06','3':'09','4':'12'}),
            format='%Y%m'
        ) + pd.offsets.QuarterEnd(0)
        
        # 筛选实际时间范围
        mask = (df['date'] >= self.start_date) & (df['date'] <= self.end_date)
        return df[mask].set_index('date')['gdp_yoy']

    def _load_macro_monthly(self) -> pd.DataFrame:
        """加载指定时间范围的PMI/PPI/CPI数据"""
        # 转换时间参数
        start_m = self._date_to_month(self.start_date)
        end_m = self._date_to_month(self.end_date)
        
        # 获取原始数据
        pmi = self.pro.cn_pmi(
            start_m=start_m, 
            end_m=end_m,
            fields='month,pmi010000'
        ).rename(columns={'pmi010000': 'pmi'})
        
        ppi = self.pro.cn_ppi(
            start_m=start_m,
            end_m=end_m,
            fields='month,ppi_yoy'
        )
        
        cpi = self.pro.cn_cpi(
            start_m=start_m,
            end_m=end_m,
            fields='month,nt_yoy'
        ).rename(columns={'nt_yoy': 'cpi'})
        
        # 统一日期处理
        for df in [pmi, ppi, cpi]:
            df['date'] = pd.to_datetime(df['month'], format='%Y%m') + pd.offsets.MonthEnd(0)
            
        # 合并并筛选
        merged = pd.merge(pmi, ppi, on='date').merge(cpi, on='date')
        mask = (merged['date'] >= self.start_date) & (merged['date'] <= self.end_date)
        return merged[mask][['date', 'pmi', 'ppi_yoy', 'cpi']].set_index('date')

    def _load_zz1000(self) -> pd.DataFrame:
        """加载指定时间范围的中证1000日线数据"""
        # 转换日期格式
        start_str = self.start_date.strftime("%Y%m%d")
        end_str = self.end_date.strftime("%Y%m%d")
        
        # 获取数据
        df = self.pro.index_daily(
            ts_code='000852.SH',
            start_date=start_str,
            end_date=end_str,
            fields='trade_date,open,high,low,close,vol'
        )
        
        # 日期处理
        df['date'] = pd.to_datetime(df['trade_date'])
        df = df.set_index('date').sort_index().asfreq('D', method='ffill')
        return df.loc[self.start_date:self.end_date, ['open', 'high', 'low', 'close', 'vol']]

    def get_market_data(self) -> dict:
        """获取指定时间范围内的全量数据"""
        return {
            'gdp': self.gdp,
            'macro': self.macro_monthly,
            'zz1000': self.zz1000
        }
